# -*- coding: utf-8 -*-
# Author   : ZhangQing
# Time     : 2025-08-04 23:35
# File     : completion_chain.py
# Project  : codebuddy_craft
# Desc     : 智能补全链

# backend/core/chains/completion_chain.py
from typing import Dict, List, Optional, Any, Tuple
import asyncio
import re
from dataclasses import dataclass

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import PydanticOutputParser
from pydantic import BaseModel, Field

from backend.utils.config import get_settings


class CompletionSuggestion(BaseModel):
    """补全建议模型"""
    text: str = Field(description="补全文本")
    label: str = Field(description="显示标签")
    detail: Optional[str] = Field(default=None, description="详细说明")
    insert_text: str = Field(description="插入文本")
    kind: str = Field(description="补全类型", default="text")
    score: float = Field(description="相关性评分", ge=0, le=1)


class CompletionResult(BaseModel):
    """补全结果模型"""
    suggestions: List[CompletionSuggestion] = Field(default=[], description="补全建议")
    context: Dict[str, Any] = Field(default={}, description="上下文信息")


@dataclass
class CodeContext:
    """代码上下文"""
    variables: List[str]
    functions: List[str]
    classes: List[str]
    imports: List[str]
    current_function: Optional[str]
    current_class: Optional[str]
    language: str


class CompletionChain:
    """智能补全链"""
    
    def __init__(self):
        self.settings = get_settings()
        # 使用正确的大写属性名
        api_key = self.settings.OPENAI_API_KEY or "sk-dummy-api-key-for-testing"
        self.llm = ChatOpenAI(
            model=self.settings.OPENAI_MODEL,
            temperature=0.1,
            api_key=api_key
        )
    
    async def get_completions(
        self,
        code: str,
        line: int,
        column: int,
        language: str = "python"
    ) -> CompletionResult:
        """获取智能补全建议"""
        try:
            # 1. 分析代码上下文
            context = self._analyze_context(code, line, column, language)
            
            # 2. 获取当前输入
            current_input = self._get_current_input(code, line, column)
            
            # 3. 生成补全建议
            suggestions = await self._generate_suggestions(
                code, current_input, context, line, column
            )
            
            return CompletionResult(
                suggestions=suggestions,
                context={
                    "variables": context.variables,
                    "functions": context.functions,
                    "classes": context.classes,
                    "current_input": current_input
                }
            )
        
        except Exception as e:
            return CompletionResult(
                suggestions=[],
                context={"error": str(e)}
            )
    
    def _analyze_context(self, code: str, line: int, column: int, language: str) -> CodeContext:
        """分析代码上下文"""
        lines = code.split('\n')
        
        variables = []
        functions = []
        classes = []
        imports = []
        current_function = None
        current_class = None
        
        if language.lower() == "python":
            # Python代码分析
            for i, code_line in enumerate(lines[:line]):
                code_line = code_line.strip()
                
                # 提取变量
                var_match = re.findall(r'^(\w+)\s*=', code_line)
                variables.extend(var_match)
                
                # 提取函数
                func_match = re.search(r'^def\s+(\w+)', code_line)
                if func_match:
                    functions.append(func_match.group(1))
                    if i < line:
                        current_function = func_match.group(1)
                
                # 提取类
                class_match = re.search(r'^class\s+(\w+)', code_line)
                if class_match:
                    classes.append(class_match.group(1))
                    if i < line:
                        current_class = class_match.group(1)
                
                # 提取导入
                import_match = re.search(r'^(?:from\s+\w+\s+)?import\s+(.+)', code_line)
                if import_match:
                    imports.extend([imp.strip() for imp in import_match.group(1).split(',')])
        
        elif language.lower() in ["javascript", "typescript"]:
            # JavaScript/TypeScript代码分析
            for i, code_line in enumerate(lines[:line]):
                code_line = code_line.strip()
                
                # 提取变量
                var_match = re.findall(r'(?:var|let|const)\s+(\w+)', code_line)
                variables.extend(var_match)
                
                # 提取函数
                func_match = re.search(r'function\s+(\w+)|(\w+)\s*=\s*(?:function|\()', code_line)
                if func_match:
                    func_name = func_match.group(1) or func_match.group(2)
                    functions.append(func_name)
                    if i < line:
                        current_function = func_name
                
                # 提取类
                class_match = re.search(r'class\s+(\w+)', code_line)
                if class_match:
                    classes.append(class_match.group(1))
                    if i < line:
                        current_class = class_match.group(1)
                
                # 提取导入
                import_match = re.search(r'import\s+.*from\s+[\'"](.+)[\'"]', code_line)
                if import_match:
                    imports.append(import_match.group(1))
        
        return CodeContext(
            variables=list(set(variables)),
            functions=list(set(functions)),
            classes=list(set(classes)),
            imports=list(set(imports)),
            current_function=current_function,
            current_class=current_class,
            language=language
        )
    
    def _get_current_input(self, code: str, line: int, column: int) -> str:
        """获取当前输入的部分文本"""
        lines = code.split('\n')
        if line > len(lines):
            return ""
        
        current_line = lines[line - 1] if line > 0 else ""
        
        # 获取光标前的文本
        prefix = current_line[:column] if column <= len(current_line) else current_line
        
        # 提取当前正在输入的词
        words = re.findall(r'\w+', prefix)
        return words[-1] if words else ""
    
    async def _generate_suggestions(
        self,
        code: str,
        current_input: str,
        context: CodeContext,
        line: int,
        column: int
    ) -> List[CompletionSuggestion]:
        """生成补全建议"""
        suggestions = []
        
        # 1. 基于上下文的简单补全
        context_suggestions = self._get_context_suggestions(current_input, context)
        suggestions.extend(context_suggestions)
        
        # 2. 使用LLM生成智能补全
        if len(current_input) >= 2:  # 输入至少2个字符才调用LLM
            llm_suggestions = await self._get_llm_suggestions(
                code, current_input, context, line, column
            )
            suggestions.extend(llm_suggestions)
        
        # 3. 去重和排序
        unique_suggestions = self._deduplicate_suggestions(suggestions)
        sorted_suggestions = sorted(unique_suggestions, key=lambda x: x.score, reverse=True)
        
        return sorted_suggestions[:10]  # 返回前10个建议
    
    def _get_context_suggestions(
        self,
        current_input: str,
        context: CodeContext
    ) -> List[CompletionSuggestion]:
        """基于上下文的简单补全"""
        suggestions = []
        
        # 变量补全
        for var in context.variables:
            if var.startswith(current_input) and var != current_input:
                suggestions.append(CompletionSuggestion(
                    text=var,
                    label=var,
                    detail="变量",
                    insert_text=var,
                    kind="variable",
                    score=0.8
                ))
        
        # 函数补全
        for func in context.functions:
            if func.startswith(current_input) and func != current_input:
                suggestions.append(CompletionSuggestion(
                    text=f"{func}()",
                    label=func,
                    detail="函数",
                    insert_text=f"{func}()",
                    kind="function",
                    score=0.9
                ))
        
        # 类补全
        for cls in context.classes:
            if cls.startswith(current_input) and cls != current_input:
                suggestions.append(CompletionSuggestion(
                    text=cls,
                    label=cls,
                    detail="类",
                    insert_text=cls,
                    kind="class",
                    score=0.7
                ))
        
        return suggestions
    
    async def _get_llm_suggestions(
        self,
        code: str,
        current_input: str,
        context: CodeContext,
        line: int,
        column: int
    ) -> List[CompletionSuggestion]:
        """使用LLM生成智能补全"""
        try:
            completion_prompt = ChatPromptTemplate.from_template("""
            作为{language}编程专家，请为以下代码提供智能补全建议：
            
            代码上下文:
            ```{language}
            {code}
            ```
            
            当前位置：第{line}行，第{column}列
            当前输入：{current_input}
            
            可用的变量：{variables}
            可用的函数：{functions}
            可用的类：{classes}
            
            请提供5个最相关的补全建议，每个建议包含：
            1. 补全文本
            2. 简短说明
            3. 使用示例（如果适当）
            
            请考虑：
            - 代码上下文的逻辑
            - 变量和函数的作用域
            - 编程最佳实践
            - 常用模式和惯例
            """)
            
            chain = completion_prompt | self.llm
            response = await chain.ainvoke({
                "language": context.language,
                "code": code,
                "line": line,
                "column": column,
                "current_input": current_input,
                "variables": ", ".join(context.variables[:10]),
                "functions": ", ".join(context.functions[:10]),
                "classes": ", ".join(context.classes[:10])
            })
            
            # 解析LLM响应
            return self._parse_llm_suggestions(response.content, current_input)
        
        except Exception as e:
            print(f"LLM补全失败: {e}")
            return []
    
    def _parse_llm_suggestions(self, response: str, current_input: str) -> List[CompletionSuggestion]:
        """解析LLM的补全建议"""
        suggestions = []
        
        # 简单解析逻辑 - 查找代码块和建议
        lines = response.split('\n')
        
        for line in lines:
            line = line.strip()
            
            # 查找可能的补全建议
            if current_input and current_input in line:
                # 提取补全文本
                parts = line.split(current_input)
                if len(parts) > 1:
                    completion = parts[1].split()[0] if parts[1].split() else ""
                    if completion:
                        full_text = current_input + completion
                        suggestions.append(CompletionSuggestion(
                            text=full_text,
                            label=full_text,
                            detail="AI建议",
                            insert_text=full_text,
                            kind="ai",
                            score=0.6
                        ))
        
        return suggestions[:5]  # 限制LLM建议数量
    
    def _deduplicate_suggestions(
        self,
        suggestions: List[CompletionSuggestion]
    ) -> List[CompletionSuggestion]:
        """去除重复建议"""
        seen = set()
        unique_suggestions = []
        
        for suggestion in suggestions:
            if suggestion.text not in seen:
                seen.add(suggestion.text)
                unique_suggestions.append(suggestion)
        
        return unique_suggestions


# 创建全局实例
completion_chain = CompletionChain()
